"""
Created on 10-Dec-2009

@author: dexen
"""

import os
import sys
import time
import traceback

import rpyc

from dexen.system.node.master import comm
from dexen.system.node.master.taskq_mgr import TaskQMgr
from dexen.system.base.spec import DexenSpec
from dexen.system.server.info import SlaveInfo
from dexen.system.base import service, task

#===============================================================================
# 
#===============================================================================
class SlaveMgr(object):
    def __init__(self, spec, tq_mgr):
        self.spec = spec
        self.tq_mgr = tq_mgr
        self.slaves_info = {}
        self.conns = {}
        self.sent_tasks = {}
        
    def add_slaves(self, slave_names):
        for slave_name in slave_names:
            print "I am creating connection to slave", slave_name
            self.slaves_info[slave_name] = SlaveInfo(slave_name)
            conn = rpyc.connect_by_service(slave_name,
                                           config=service.default_config)
            conn.root.run(self.spec)
            self.conns[slave_name] = conn
            self.sent_tasks[slave_name] = []
    
    # TODO: cater for crashes
    def _update(self):
        #print "updating slaves"
        for slave_name, conn in self.conns.items():
            if self.slaves_info[slave_name].is_crashed():
                continue
            try:
                outputs = conn.root.get()
                outputs = rpyc.classic.obtain(outputs)
            except EOFError:
                self._slave_unreachable(slave_name)
                self.slaves_info[slave_name].crashed()
            else:
                if len(outputs) > 0:
                    print "I have collected", len(outputs), "output from", slave_name
                for unit in outputs:
                    assert isinstance(unit, task.OutputUnit)
                    self.tq_mgr.set_inds(unit.task_name, unit.output)
                    self.sent_tasks[slave_name] = self.sent_tasks[slave_name][1:]
    
    def _recover(self, pending_tasks):
        for task_name, inds in pending_tasks:
            self.tq_mgr.recover_inds(task_name, inds)
    
    def _slave_unreachable(self, slave_name):
        print "I am recovering for slave:", slave_name 
        self._recover(self.sent_tasks[slave_name])
        self.sent_tasks[slave_name] = [] # TODO: beware
        
    
    def _get_available_slave(self):
        #print "finding available slave"
        for slave_name, conn in self.conns.items():
            if self.slaves_info[slave_name].is_crashed():
                continue
            try:
                if conn.root.available():
                    return slave_name
            except EOFError:
                self._slave_unreachable(slave_name)
        return None
    
    def execute(self):
        if len(self.slaves_info) <= 0:
            return    
        
        self._update()
        
        slave_name = self._get_available_slave()
        if not slave_name:
            return
        conn = self.conns[slave_name]
    
        task_name, inds = self.tq_mgr.get_inds()
        if task_name is None: # means none of the tasks have enough input data
            return
        
        print "I am putting", task_name, len(inds), "into", slave_name
        print self.tq_mgr
        
        try:
            conn.root.put(task_name, inds) # Might need id to match the tasks!
            self.sent_tasks[slave_name].append((task_name, inds))
        except EOFError: # TODO: check the exception type (slave is dead)
            self.tq_mgr.recover_inds(task_name, inds)
        
        # Make sure that the slave side never returns an exception when alive

        #print conn.get_input()
    
    def get_info(self):
         return self.slaves_info

def test(queue):
    try:
        print queue.get()
    except Exception:
        traceback.print_exc()


